// Assumptions: Supports only one client at a time
// client == ALICE

#include "ramprakash.Part2Common.h"

int servSockFd;
fd_set writeFds;
struct sockaddr_in servaddr;
BIGNUM *gP;// large prime number (p)
BIGNUM *gG;// Generator for gP
BIGNUM *gXA;// Alice's secret - random number (>p)
BIGNUM *gKAS; // Session Key
BIGNUM *gNB; // Bob's Nonce
BIGNUM *gNA; // Alice's Nonce
char gDesKey[8];


int connectToServer(char *ip)
{
    int buflen = 2048;
    servSockFd = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (servSockFd < 0) 
    {
        perror("socket:");
        return errno;
    }

    setsockopt(servSockFd, SOL_SOCKET, SO_SNDBUF, &buflen, sizeof(int));
    setsockopt(servSockFd, SOL_SOCKET, SO_RCVBUF, &buflen, sizeof(int));

    memset(&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_port = htons(LISTENING_PORT);
    inet_aton(ip, &servaddr.sin_addr);

    if (connect(servSockFd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
    {
        perror("connect");
        return errno;
    }

    return 0;
}

void generateAuthenticationParams()
{
    const char *pp = "197221152031991558322935568090317202983";
    const char *two = "2";
    int         rand = 0;
    char        rand_buf[10];
#ifdef DEBUG
    char       *bn_buf;
    int         llen = 0;
#endif

    // Generating gP and gG
    if (BN_dec2bn(&gP, pp) == 0) {
        printf("BN_dec2bn failed\n");
    }
    assert(gP);
    if (BN_dec2bn(&gG, two) == 0) {
        printf("BN_dec2bn failed\n");
        exit(1);
    }
    assert(gG);
#ifdef DEBUG
    llen = BN_num_bytes(gP);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gP);
    printf("\ngP: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif

    rand = getRandom();
    memset(rand_buf, 0, sizeof(rand_buf));
    sprintf(rand_buf, "%x", rand);

    // Alice's context:
    // Generate gXA - ensure it is greater than gP by generating random
    // number and adding it to gP
    if (BN_hex2bn(&gXA, rand_buf) == 0) {
        printf("BN_hex2bn failed\n");
    }
    if (gXA == NULL) {
        printf("BN_hex2bn failed : gXA is NULL!\n");
        exit(1);
    }
    // gXA = gXA + gP
    BN_add(gXA, gP, gXA);
#ifdef DEBUG
    llen = BN_num_bytes(gXA);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gXA);
    printf("gXA: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif
}

int handleMsg4(msg_t *pMsg, char *pw)
{
    BIGNUM *bn = NULL;
    char pt[200];
    memset(pt, 0, sizeof(pt));
    desCryptM(gDesKey, pMsg->u.m4.enc_des_na, pMsg->u.m4.len_des_na, pt, 0);
    BN_hex2bn(&bn, pt);

    if (BN_cmp(bn, gNA) == 0) {
        LOG("Authentication succeeded - proceeding with file retrieval\n");
        return 0;
    } 
    return 1;
}
void handleMsg2(msg_t *pMsg, char *pw)
{
    char pt[200];
    char *buf;
#ifdef DEBUG
    char *bn_buf;
    int llen;
#endif

    memset(pt, 0, sizeof(pt));
    cryptM(pw, pMsg->u.m2.enc_gkas_modp, pMsg->u.m2.len_gkas_modp, pt, 0);
    BN_hex2bn(&gKAS, pt);
    buf = BN_bn2hex(gKAS);
    memcpy(gDesKey, buf, 8);
#ifdef DEBUG
    llen = BN_num_bytes(gKAS);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gKAS);
    printf("gKAS: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif

    memset(pt, 0, sizeof(pt));
    desCryptM(gDesKey, pMsg->u.m2.enc_des_kas, pMsg->u.m2.len_des_kas, pt, 0);
    BN_hex2bn(&gNB, pt);
#ifdef DEBUG
    llen = BN_num_bytes(gNB);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gNB);
    printf("gNB: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif
}

void sendMessage3()
{
#ifdef DEBUG
    char *bn_buf;
#endif
    msg_t msg;
    char        rand_buf[10];
    int   rand;
    int   len;
    int   deslen;
    char *buf;

    memset (&msg, 0, sizeof(msg));
    msg.cmd = MSG_3;

    // Generate Alice's nonce (gNA) - a random int more than Bob's nonce (gNB)
    rand = getRandom();
    memset(rand_buf, 0, sizeof(rand_buf));
    sprintf(rand_buf, "%x", rand);
    if (BN_hex2bn(&gNA, rand_buf) == 0) {
        printf("BN_hex2bn failed\n");
        exit(1);
    }
    BN_add(gNA, gNB, gNA);
#ifdef DEBUG
    len = BN_num_bytes(gNA);
    bn_buf = malloc(len);
    bn_buf = BN_bn2hex(gNA);
    printf("gNA: %s of length: %d\n", bn_buf, len);
    free(bn_buf);
#endif

    //Encrypt using DES with gKAS(first 8 bytes) as key and 
    //encrypt the Nonce gNA
    len = BN_num_bytes(gNA);
    len *= 2;
    buf = BN_bn2hex(gNA);
    deslen = desCryptM(gDesKey, buf, len, msg.u.m3.enc_des_na, 1);
    msg.u.m3.len_des_na = deslen;

    //Encrypt using DES with gKAS(first 8 bytes) as key and 
    //encrypt the Nonce gNB
    len = BN_num_bytes(gNB);
    len *= 2;
    buf = BN_bn2hex(gNB);
    deslen = desCryptM(gDesKey, buf, len, msg.u.m3.enc_des_nb, 1);
    msg.u.m3.len_des_nb = deslen;

    sendMsg(servSockFd, (char *)&msg, sizeof(msg));
}

void sendAuthReq(char *u, char *p)
{
    msg_t msg;

    memset(&msg, 0, sizeof(msg));
    msg.cmd = AUTH_REQ;
    desCryptM(gDesKey, u, 32, msg.u.a.user, 1);
    desCryptM(gDesKey, p, 64, msg.u.a.passwd, 1);
    sendMsg(servSockFd, (char *)&msg, sizeof(msg));
}

// 3 tries at getting authenticated to the server with plaintext user/password
int authenticateUser(char *u, char *p, int sockFd, struct sockaddr_in *srv_src)
{
    fd_set readFds;
    struct timeval tv;
    struct sockaddr_in src;
    int rc;
    int auth_try = 0;
    msg_t msg, *pMsg;
    socklen_t len = 0;
    char lu[32];
    char lp[64];
    char message[1500];
    char encrypted_m[1500];
    BN_CTX *ctx;
    BIGNUM *bn;
    int llen;

    memset(&message, 0, sizeof(message));
    memset(&lu, 0, sizeof(lu));
    memset(&lp, 0, sizeof(lp));
    strcpy(lu, u);
    strcpy(lp, p);

    generateAuthenticationParams();

next_try:

    auth_try++;
    memset(&src, 0, sizeof(src));
    memset(&msg, 0, sizeof(msg));

    msg.cmd = MSG_1;
    strcpy(msg.u.m1.user, lu);

    // Calculate (g^Xa mod p)
    ctx = BN_CTX_new();
    bn  = BN_new();
    BN_mod_exp(bn, gG, gXA, gP, ctx);
#ifdef DEBUG
    char *bn_buf;
    llen = BN_num_bytes(bn);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(bn);
    printf("gGXA: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif
    // Get E(pass, g^Xa mod p) in encrypted_m
    memset(&encrypted_m, 0, sizeof(encrypted_m));
    llen = BN_num_bytes(bn);
    llen *= 2;
    cryptM(lp, BN_bn2hex(bn), llen, encrypted_m, 1);
    msg.u.m1.len_gxa_modp = llen;
    memcpy(msg.u.m1.enc_gxa_modp, encrypted_m, llen);

    memset(&encrypted_m, 0, sizeof(encrypted_m));
    llen = BN_num_bytes(gP);
    llen *= 2;
    cryptM(lp, BN_bn2hex(gP), llen, encrypted_m, 1);
    msg.u.m1.len_p = llen;
    memcpy(msg.u.m1.enc_p, encrypted_m, llen);

    FD_ZERO(&readFds);
    tv.tv_sec = 2;
    tv.tv_usec = 0;

    if (sendto(sockFd, &msg, sizeof(msg), 0, (const struct sockaddr *)srv_src, 
               srv_src ?  sizeof(struct sockaddr_in) : 0) == -1) 
    {
        perror("sendto:");
        return errno;
    }

waitforpkt:

    FD_SET(sockFd, &readFds);
    rc = select(sockFd+1, &readFds, NULL, NULL, &tv);
    if (rc == 0) {
        printf("No response from the server. Check if server is running!\n");
        return -1;
    }

    memset(&msg, 0, sizeof(msg));
    if ((len = recvfrom(sockFd, &msg, sizeof(msg), 0, 
                    (struct sockaddr *)&src, 
                    &len)) > 0) {
        pMsg = &msg;
    } else {
        perror("recvfrom");
        return errno;
    }

    if (pMsg->cmd == MSG_2) {
        handleMsg2(pMsg, p);
        sendMessage3();
        goto waitforpkt;
    }

    if (pMsg->cmd == MSG_4) {
        if (handleMsg4(pMsg, p) != 0) {
            printf("Shared Key mismatch (Na != Nb)\n");
            if (getPass(lu, lp) != 0)
                exit(-1);
            goto next_try;
        }
        sendAuthReq(lu, lp);
        goto waitforpkt;
    }

    if (pMsg->cmd == AUTH_RESP) {
        if (strcmp(pMsg->u.a.user, "Success") == 0) {
            LOG("Authentication: Success\n");
            return 0;
        }
    }
    if (auth_try < 3) {
        printf("\nAuthentication Failed\n");
        if (getPass(lu, lp) != 0)
            exit(-1);

        goto next_try;
    }
    exit(-1);
}

int getFile(char *inrf, char *inlf, int sockFd, 
            struct sockaddr_in *srv_src, unsigned char *md)
{
    struct sockaddr_in src;
    msg_t msg;
    file_t f;
    int len;
    socklen_t slen = 0;
    FILE *lfd;
    char rf[80], lf[80];

    memset(rf, 0, sizeof(rf));
    memset(lf, 0, sizeof(lf));
    strcpy(lf, inlf);
    strcpy(rf, inrf);

next_try:
    memset(&src, 0, sizeof(src));
    memset(&msg, 0, sizeof(msg));
    msg.cmd = FILE_REQ;
    //strcpy(msg.u.f.file, rf);
    desCryptM(gDesKey, rf, 80, msg.u.f.file, 1);

    if (sendto(sockFd, &msg, sizeof(msg), 0, (const struct sockaddr *)srv_src,
               srv_src ? sizeof(struct sockaddr_in) : 0) == -1) {
        perror("sendto:");
        return errno;
    }

    lfd = fopen(lf, "w+");
    if (!lfd) {
        perror("fopen");
        return errno;
    }

    memset(&msg, 0, sizeof(msg));
    do {
        memset(&f, 0, sizeof(f));
        if ((len = recvfrom(sockFd, (char *)&msg, sizeof(msg_t), 0, 
                            (struct sockaddr *)&src, 
                            &slen)) > 0) {
            while (len != sizeof(msg_t)) {
                len += recvfrom(sockFd, (char *)&msg + len, 
                                sizeof(msg_t) - len, 0, 
                                (struct sockaddr *)&src, &slen);
            }
            desCryptM(gDesKey, (char *)&msg.u.f, msg.len, (char *)&f, 0);
        } else {
            perror("recvfrom");
            return errno;
        }

        if ((msg.cmd == FILE_RESP) && (f.size == 0)) {
            printf("Remote file not found. Retry with remote file name:");
            scanf("%s", rf);
            goto next_try;
        }

        if (fwrite(f.data, 1, f.chunk_size, lfd) < f.chunk_size) {
            perror("fwrite");
            return errno;
        }
    } while ((f.chunk_size + f.offset) < f.size);
    fclose(lfd);
    memcpy(md, f.digest, 20);

    return 0;
}


int main(int argc, char *argv[])
{
    char u[32], p[64];
    unsigned char md[20];
    unsigned char mymd[20];
    int j;
#ifdef DESDEBUG
    // This piece of code is just to test the SHA1 functionalisty. 
    // run tcpClnt with filename for which SHA1 needs to be calculated.
    // To compile with this flag set, run "make desdebug"
    // The digest generated can be verified by running the command 
    // "shasum <filename>" and checking if the same digest is generated.
    int i;

    if (argc != 2) {
        printf("Usage: %s <filename>\n", argv[0]);

        return 0;
    }
    sha1Digest(argv[1], md);
    printf("SHA-1 digest: ");
    for(i = 0; i < 20; i++) {
        printf("%02x", md[i]);
    }
    printf("\n");
    return 0;
#endif

    if (argc != 4) {
        printf("Usage: %s <server-ip-addr> <full-path-to-remote-file> "
               "<local-file>\n", argv[0]);
        return -1;
    }

    if (connectToServer(argv[1]) != 0) {
        return -1;
    }

    memset(u, 0, sizeof(u));
    memset(p, 0, sizeof(p));
    if (getPass(u, p) != 0) {
        return -1;
    }
    if (authenticateUser(u, p, servSockFd, NULL) < 0) {
        printf("Authentication failed\n");
        return -1;
    }

    if (getFile(argv[2], argv[3], servSockFd, NULL, md) < 0) {
        printf("Failed to get file\n");
        return -1;
    }

    printf("Retrieved file %s as %s\n\n", argv[2], argv[3]);
    memset(mymd, 0, 20);
    sha1Digest(argv[3], mymd);
    if (memcmp(md, mymd, 20) == 0) {
        printf("SHA-1 Digest received matches the calculated SHA-1 Digest\n");
        printf("Digest: ");
        for (j = 0; j < 20; j++) {
            printf("%02x", md[j]);
        }
        printf("\n");
    } else {
        printf("SHA-1 Digest mismatch. Check if Trudy is online!!\n");
        printf("The SHA-1 digest received: ");
        for (j = 0; j < 20; j++) {
            printf("%02x", md[j]);
        }
        printf("\n");

        printf("The calculated digest    : ");
        for (j = 0; j < 20; j++) {
            printf("%02x", mymd[j]);
        }
        printf("\n");
    }

    return (0);
}
